import os
import signal
import threading
import asyncio
import websockets

connected_clients = set()


async def handle_client(websocket, path):
    global connected_clients
    connected_clients.add(websocket)
    try:
        while True:
            message = await websocket.recv()
            # 如果recv()抛出了ConnectionClosed异常，将会被下面的except捕获
            # 广播消息给所有客户端
            for client in connected_clients:
                await client.send(message)
    except websockets.exceptions.ConnectionClosed as e:
        print(f"客户端 {websocket} 断开连接，状态码：{e.code}，原因：{e.reason}")
    except Exception as e:
        print(f"客户端 {websocket} 断开连接，异常：{e}")
    finally:
        # 客户端断开连接时移除
        connected_clients.remove(websocket)


def main():
    start_server = websockets.serve(handle_client, "localhost", 8899)

    asyncio.get_event_loop().run_until_complete(start_server)
    asyncio.get_event_loop().run_forever()


def input_cmd():
    help_str = """
1.显示当前连接的客户端
#2.主动断开 指定客户端
3.断开所有 客户端
4.kill SIGINT 自己
q.退出
"""
    global connected_clients
    while True:
        print(help_str)
        cmd = input("请输入指令：")
        if cmd == "1":
            print("当前连接的客户端：", connected_clients)
        # elif cmd == "2":
        #     idx = input("请输入要断开的客户端：")
        #     try:
        #         client = connected_clients[idx]
        #         client.close()
        #         connected_clients.remove(client)
        #         print("已断开客户端：", client)
        #     except:
        #         print("客户端不存在")
        elif cmd == "3":
            for client in connected_clients:
                client.close()
            connected_clients.clear()
            print("已断开所有客户端")
        elif cmd == "4":
            os.kill(os.getpid(), signal.SIGINT)
        elif cmd == "q":
            break
        else:
            print("无效指令")


if __name__ == "__main__":
    t1 = threading.Thread(target=input_cmd)
    t1.daemon = True
    t1.start()

    main()
